#ifndef __TILING_COMMON_H_
#define __TILING_COMMON_H_
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <cstdint>
#include <iostream>
#define COMMON_TILING_FILED_DEF                                                                                        \
    TILING_DATA_FIELD_DEF(int64_t, size);                                                                             \
    TILING_DATA_FIELD_DEF(int64_t, formerNum);                                                                         \
    TILING_DATA_FIELD_DEF(int64_t, formerLength);                                                                      \
    TILING_DATA_FIELD_DEF(int64_t, formerTileLength);                                                                  \
    TILING_DATA_FIELD_DEF(int64_t, formerFinalCalcCount);                                                              \
    TILING_DATA_FIELD_DEF(int64_t, tailNum);                                                                           \
    TILING_DATA_FIELD_DEF(int64_t, tailLength);                                                                        \
    TILING_DATA_FIELD_DEF(int64_t, tailTileLength);                                                                    \
    TILING_DATA_FIELD_DEF(int64_t, tailFinalCalcCount);                                                                \
    TILING_DATA_FIELD_DEF(int64_t, dType);

#define COMMON_TILING_TAG_DEF TILING_DATA_FIELD_DEF(uint64_t, tag)

#define COMMON_TILING_SPLIT(tiling, blockDim, alignCount, alignDataSize)                                               \
    auto kernelBlock = alignCount / blockDim;                                                                          \
    auto formerNum = alignCount % blockDim;                                                                            \
    auto formerLength = (kernelBlock + 1) * alignDataSize;                                                             \
    tiling->set_formerNum(formerNum);                                                                                  \
    tiling->set_formerLength(formerLength);                                                                            \
    auto tailNum = blockDim - formerNum;                                                                               \
    auto tailLength = kernelBlock * alignDataSize;                                                                     \
    tiling->set_tailNum(tailNum);                                                                                      \
    tiling->set_tailLength(tailLength);

#define CEIL_DIV(a, b) ((a + b - 1) / b)
#define ALIGN_TO(a, b) CEIL_DIV(a, b) * b
namespace optiling
{
typedef int64_t (*mem_addition_fn)(ge::DataType);
template <typename _T_TILING> void print_common_tiling(_T_TILING &tiling)
{
    std::cout << "get_size : " << tiling.get_size() << std::endl;
    std::cout << "get_formerNum : " << tiling.get_formerNum() << std::endl;
    std::cout << "get_formerLength : " << tiling.get_formerLength() << std::endl;
    std::cout << "get_formerTileLength : " << tiling.get_formerTileLength() << std::endl;
    std::cout << "get_formerFinalCalcCount : " << tiling.get_formerFinalCalcCount() << std::endl;
    std::cout << "get_tailNum : " << tiling.get_tailNum() << std::endl;
    std::cout << "get_tailLength : " << tiling.get_tailLength() << std::endl;
    std::cout << "get_tailTileLength : " << tiling.get_tailTileLength() << std::endl;
    std::cout << "get_tailTileLength : " << tiling.get_tailTileLength() << std::endl;
    std::cout << "get_tailFinalCalcCount : " << tiling.get_tailFinalCalcCount() << std::endl;
    std::cout << "get_dType : " << tiling.get_dType() << std::endl;
}

template <typename _T_TILING> void tiling_save_to_buffer(_T_TILING &tiling, gert::TilingContext *context)
{
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
}

template <typename _T_TILING>
void set_size_by_input(_T_TILING &tiling, gert::TilingContext *context, int64_t shape_index)
{
    auto shape = context->GetInputShape(shape_index);
    auto dataSize = shape->GetStorageShape().GetShapeSize();
    tiling.set_size(dataSize);
}

void tiling_zero_workspace(gert::TilingContext *context)
{
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
}

template <typename _T_TILING> int64_t common_tiling_set_kernel_vec(_T_TILING &tiling, gert::TilingContext *context)
{
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto aicNum = ascendcPlatform.GetCoreNumAic();
    auto aivNum = ascendcPlatform.GetCoreNumAiv();
    auto blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum);
    context->SetBlockDim(blockDim);
    return blockDim;
}

int64_t common_tiling_auto_tileLenth(platform_ascendc::PlatformAscendC &ascendcPlatform, int64_t dtypeSize,
                                     int64_t alignDataSize, int64_t total_count, int64_t alignedSize)
{
    // 计算UB可以容纳多少个数据
    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    // 计算每个核能容纳多少个操作数
    ub_size /= dtypeSize;
    // 由于非对齐shape的原因，最后一次循环可能会超出，所以为其配置容量
    ub_size -= alignDataSize;
    // ub_size / total_count 然后对齐 alignDataSize
    auto ret = ub_size / total_count / alignDataSize * alignDataSize;
    if (alignedSize < ret)
    {
        return alignedSize;
    }
    return ret;
}

/*
tiling：TilingData
context: gert::TilingContext
vector_align_size: 对齐大小 默认32, 单位：Byte。
tile_align_size: 一次搬运多少个 vector_align_size 的数据。 默认 8
input_arg_count: 算子的输入个数。 默认:1
output_arg_count: 算子的输出个数。默认:1
auto_tilelength: 是否自动计算 tile_align_size， 如果为true, tile_align_size参数将失效。默认：true
mem_addition： 好需要多少个 tile_align_size 的 空间。一般用来存储Tbuf。 默认：nullptr
target_(dtype/shape)_(input/output)_index: 使用(input/output)的(dtype/shape)来进行切分。 -1 表示不用。默认 input 0
*/
template <typename _T_TILING> class CommonTiling
{
  public:
    _T_TILING *tiling;
    gert::TilingContext *context = nullptr;
    int64_t data_size;
    uint64_t vector_align_size = 32;
    uint64_t tile_align_size = 8;
    uint64_t input_arg_count = 1;
    uint64_t output_arg_count = 1;
    bool auto_tilelength = true;
    mem_addition_fn mem_addition = nullptr;
    int64_t mem_addition_num = 0;

  public:
    CommonTiling()
    {
    }
    ~CommonTiling()
    {
    }

    void apply()
    {
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        auto aicNum = ascendcPlatform.GetCoreNumAic();
        auto aivNum = ascendcPlatform.GetCoreNumAiv();
        auto blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum);
        context->SetBlockDim(blockDim);

        // 32B对齐
        auto dtype = context->GetInputDesc(0)->GetDataType();
        auto dtypeSize = GetSizeByDataType(dtype);
        auto alignDataSize = vector_align_size / dtypeSize;
        auto alignCount = (data_size + alignDataSize - 1) / alignDataSize;
        auto alignedSize = alignCount * alignDataSize;

        // 分割
        COMMON_TILING_SPLIT(tiling, blockDim, alignCount, alignDataSize)

        // 只有对于搬运耗时较大的运算才使用auto_tilelength
        int64_t formerTileLength = 0;
        int64_t tailTileLength = 0;
        if (auto_tilelength)
        {
            auto addition = mem_addition_num;
            if (mem_addition)
            {
                addition += mem_addition(dtype);
            }
            auto total_count = (input_arg_count + output_arg_count) * 2 + addition;
            formerTileLength =
                common_tiling_auto_tileLenth(ascendcPlatform, dtypeSize, alignDataSize, total_count, formerLength);
            tailTileLength =
                common_tiling_auto_tileLenth(ascendcPlatform, dtypeSize, alignDataSize, total_count, tailLength);
        }
        else
        {
            auto tileLength = alignDataSize * tile_align_size;
            formerTileLength = formerLength < tileLength ? formerLength : tileLength;
            tailTileLength = tailLength < tileLength ? tailLength : tileLength;
        }

        tiling->set_formerTileLength(formerTileLength);
        tiling->set_tailTileLength(tailTileLength);

        // 最后一次循环
        int64_t formerFinalCalcCount = ALIGN_TO((formerLength % formerTileLength), alignDataSize);
        int64_t tailFinalCalcCount = ALIGN_TO((tailLength % tailTileLength), alignDataSize);
        if (formerFinalCalcCount == 0)
        {
            formerFinalCalcCount = formerTileLength;
        }
        if (tailFinalCalcCount == 0)
        {
            tailFinalCalcCount = tailTileLength;
        }

        tiling->set_formerFinalCalcCount(formerFinalCalcCount);
        tiling->set_tailFinalCalcCount(tailFinalCalcCount);

        // 设置dtype
        tiling->set_dType(dtype);
    }
};

}; // namespace optiling
#endif